package aggregate

import (
	"context"
	"fmt"
	"github.com/kataras/golog"
	"go.uber.org/ratelimit"
	"math"
	"nprobe/core"
	"nprobe/structs"
	"nprobe/utils/routines"
	"sort"
	"time"
)

type RingBufAggregator struct {
	Ctx           context.Context
	Name          string
	Task          *core.ProbeTask
	FirstSpanID   int64
	SenderRBufs   []*structs.PPRingBuffer[*core.SpanRecord]
	ReceiverRBufs []*structs.PPRingBuffer[*core.SpanRecord]
}

func NewRingBufAggregator(
	ctx context.Context,
	name string,
	task *core.ProbeTask,
	senderGroup *core.SenderGroup,
	receiverGroup *core.ReceiverGroup,
) *RingBufAggregator {
	var senderRBufs []*structs.PPRingBuffer[*core.SpanRecord]
	if task.Role == "client" {
		for _, sender := range senderGroup.Senders {
			senderRBufs = append(senderRBufs, sender.GetRBuf())
		}
	}

	var receiverRBufs []*structs.PPRingBuffer[*core.SpanRecord]
	for _, receiver := range receiverGroup.Receivers {
		receiverRBufs = append(receiverRBufs, receiver.GetRBuf())
	}

	return &RingBufAggregator{
		Ctx:           ctx,
		Name:          name + "-aggregator",
		Task:          task,
		SenderRBufs:   senderRBufs,
		ReceiverRBufs: receiverRBufs,
	}
}

func (a *RingBufAggregator) StartAggregator() error {
	routines.GoWithRecover(a.Name, a.doAggregate)
	return nil
}

func (a *RingBufAggregator) StopAggregator() error {
	return nil
}
func (a *RingBufAggregator) doAggregate() {
	limiter := ratelimit.New(1, ratelimit.Per(time.Duration(a.Task.SpanMs)*time.Millisecond))
	for {

		select {
		case <-a.Ctx.Done():
			a.StopAggregator()
			return
		default:
			limiter.Take()
			a.tryAggregate()
		}
	}
}

func (a *RingBufAggregator) tryAggregate() {
	spanRecords := a.tryAggrregateSpanRecord()
	for _, spanRecord := range spanRecords {
		if a.FirstSpanID == 0 && spanRecord.SpanID != 0 {
			a.FirstSpanID = spanRecord.SpanID
		}
		aggregateSpanRecordAndPrint(a.FirstSpanID, spanRecord, a.Task.Role)
	}
}

func (a *RingBufAggregator) tryAggrregateSpanRecord() []*core.SpanRecord {
	var spanRecords []*core.SpanRecord
	for {
		mergedSpanRecord := core.NewSpanRecord(0)
		if a.Task.Role == "client" {
			a.mergeSenderSpanRecord(mergedSpanRecord)
			if mergedSpanRecord.SpanID == 0 {
				break
			}
			a.mergeReceiverSpanRecord(mergedSpanRecord)
		} else {
			a.mergeServerSpanRecord(mergedSpanRecord)
			if mergedSpanRecord.SpanID == 0 {
				break
			}
		}
		spanRecords = append(spanRecords, mergedSpanRecord)
	}
	return spanRecords
}

func (a *RingBufAggregator) mergeSenderSpanRecord(mergedSpanRecord *core.SpanRecord) {
	a.mergeFirstSpanRecord(mergedSpanRecord, a.SenderRBufs)
}

func (a *RingBufAggregator) mergeServerSpanRecord(mergedSpanRecord *core.SpanRecord) {
	a.mergeFirstSpanRecord(mergedSpanRecord, a.ReceiverRBufs)
}

func (a *RingBufAggregator) mergeFirstSpanRecord(mergedSpanRecord *core.SpanRecord, rbufs []*structs.PPRingBuffer[*core.SpanRecord]) {
	for _, buf := range rbufs {
		r, err := buf.PeekPong()
		if err != nil {
			continue
		}

		if mergedSpanRecord.SpanID == 0 || mergedSpanRecord.SpanID == r.SpanID {
			mergedSpanRecord.SpanID = r.SpanID
			for k, v := range r.PktRecordMap {
				mergedSpanRecord.PktRecordMap[k] = v
			}
			buf.DequeueAfterPeek()
		} else if mergedSpanRecord.SpanID > r.SpanID {
			buf.DequeueAfterPeek()
			for {
				r, err = buf.PeekPong()
				if err != nil {
					break
				}

				if mergedSpanRecord.SpanID <= r.SpanID {
					if mergedSpanRecord.SpanID == r.SpanID {
						for k, v := range r.PktRecordMap {
							mergedSpanRecord.PktRecordMap[k] = v
						}
						buf.DequeueAfterPeek()
					}
					break
				}

				buf.DequeueAfterPeek()
			}
		}
	}
}

func (a *RingBufAggregator) mergeReceiverSpanRecord(mergedSpanRecord *core.SpanRecord) {
	for _, buf := range a.ReceiverRBufs {
		r, err := buf.PeekPong()
		if err != nil {
			continue
		}

		if mergedSpanRecord.SpanID == r.SpanID {
			for k, v := range r.PktRecordMap {
				sv := mergedSpanRecord.PktRecordMap[k]
				if sv != nil {
					sv.SrcIP = v.SrcIP
					sv.Received = v.Received
					sv.BitFlip = v.BitFlip
					sv.RxTs = v.RxTs
					sv.RxTxTs = v.RxTxTs
					sv.Latency = v.RxTs - sv.TxTs
				}
			}
			buf.DequeueAfterPeek()
		} else if mergedSpanRecord.SpanID > r.SpanID {
			buf.DequeueAfterPeek()
			for {
				r, err = buf.PeekPong()
				if err != nil {
					break
				}

				if mergedSpanRecord.SpanID <= r.SpanID {
					if mergedSpanRecord.SpanID == r.SpanID {
						for k, v := range r.PktRecordMap {
							sv := mergedSpanRecord.PktRecordMap[k]
							if sv != nil {
								sv.SrcIP = v.SrcIP
								sv.Received = v.Received
								sv.BitFlip = v.BitFlip
								sv.RxTs = v.RxTs
								sv.RxTxTs = v.RxTxTs
								sv.Latency = v.RxTs - sv.TxTs
							}
						}
						buf.DequeueAfterPeek()
					}
					break
				}

				buf.DequeueAfterPeek()
			}
		}
	}
}

func aggregateSpanRecordAndPrint(firstSpanID int64, spanRecord *core.SpanRecord, role string) {
	if spanRecord != nil {
		aggreSpanRecord := aggregateSpanRecord(spanRecord)
		printAggreSpanRecord(firstSpanID, aggreSpanRecord, role)
	}
}

func aggregateSpanRecord(spanRecord *core.SpanRecord) *core.AggreSpanRecord {
	aggreSpanRecord := &core.AggreSpanRecord{
		SpanID:         spanRecord.SpanID,
		AggreRecordMap: make(map[string]*core.AggreRecord),
	}

	for _, r := range spanRecord.PktRecordMap {
		key := fmt.Sprintf("%s", r.DstIP)
		aggreRecord := aggreSpanRecord.AggreRecordMap[key]
		if aggreRecord == nil {
			aggreRecord = &core.AggreRecord{
				Key:        key,
				DstIP:      r.DstIP,
				MinLatency: math.MaxInt64,
			}
			aggreSpanRecord.AggreRecordMap[key] = aggreRecord
		}

		preTxCnt := aggreRecord.TxPktCnt
		aggreRecord.TxPktCnt++
		if !r.Received {
			aggreRecord.LossPktCnt++
		}

		if r.BitFlip {
			aggreRecord.BitFlipCnt++
		}

		if r.Latency > aggreRecord.MaxLatency {
			aggreRecord.MaxLatency = r.Latency
		}

		if r.Latency < aggreRecord.MinLatency {
			aggreRecord.MinLatency = r.Latency
		}

		aggreRecord.AvgLatency = (aggreRecord.AvgLatency*int64(preTxCnt) + r.Latency) / int64(aggreRecord.TxPktCnt)
		aggreRecord.PktRecords = append(aggreRecord.PktRecords, r)
	}

	return aggreSpanRecord
}

func printAggreSpanRecord(firstSpanID int64, r *core.AggreSpanRecord, role string) {

	currTime := time.Now().UnixNano()
	currSpanID := currTime / (1000 * int64(time.Millisecond))
	golog.Debugf("spanID: %d, currSpanID: %d, target count: %d", r.SpanID, currSpanID, len(r.AggreRecordMap))
	seq := r.SpanID - firstSpanID + 1
	if seq == 1 {
		if role == "client" {
			fmt.Printf("%-8s %-15s %-8s %-8s %-8s %-8s %-8s %-8s %-8s\n",
				"seq", "target", "total", "loss", "bitflip", "success", "min(ms)", "avg(ms)", "max(ms)")
		} else {
			fmt.Printf("%-8s %-15s %-8s %-8s %-8s %-8s %-8s %-8s %-8s\n",
				"seq", "source", "total", "loss", "bitflip", "success", "min(ms)", "avg(ms)", "max(ms)")
		}
	}

	// 打印统计信息
	printSpanRecordStat(seq, r)

	// 打印异常流
	printSpanRecordAbnormalFlow(r)
}

func printSpanRecordStat(seq int64, r *core.AggreSpanRecord) {
	var keys []string
	for k := range r.AggreRecordMap {
		keys = append(keys, k)
	}
	sort.Strings(keys)

	var i int
	for _, k := range keys {
		v := r.AggreRecordMap[k]
		minLat := float64(v.MinLatency) / float64(time.Millisecond)
		avgLat := float64(v.AvgLatency) / float64(time.Millisecond)
		maxLat := float64(v.MaxLatency) / float64(time.Millisecond)
		var rateStr string
		failPktCnt := v.LossPktCnt + v.BitFlipCnt
		if failPktCnt == 0 {
			rateStr = "100%"
		} else if failPktCnt == v.TxPktCnt {
			rateStr = "0%"
		} else {
			rate := (1 - float64(failPktCnt)/float64(v.TxPktCnt)) * 100
			rateStr = fmt.Sprintf("%.4f%%", rate)
		}

		if i == 0 {
			fmt.Printf("%-8d %-15s %-8d %-8d %-8d %-8s %-8.2f %-8.2f %-8.2f\n",
				seq, k, v.TxPktCnt, v.LossPktCnt, v.BitFlipCnt, rateStr, minLat, avgLat, maxLat)
		} else {
			fmt.Printf("%-8s %-15s %-8d %-8d %-8d %-8s %-8.2f %-8.2f %-8.2f\n",
				"", k, v.TxPktCnt, v.LossPktCnt, v.BitFlipCnt, rateStr, minLat, avgLat, maxLat)
		}
		i++
	}
}

func printSpanRecordAbnormalFlow(r *core.AggreSpanRecord) {

	keys, lossFlowMap, bitFlipFlowMap := extractAbnormalFlow(r)
	sort.Strings(keys)

	if len(keys) > 0 {
		fmt.Println()
		fmt.Println("=====================================")
	}

	if len(lossFlowMap) > 0 {
		fmt.Println("abnormal loss flows:")
		for _, key := range keys {
			for _, text := range lossFlowMap[key] {
				fmt.Println("    " + text)
			}
		}
	}

	if len(bitFlipFlowMap) > 0 {
		fmt.Println("abnormal bit flip flows:")
		for _, key := range keys {
			for _, text := range bitFlipFlowMap[key] {
				fmt.Println("    " + text)
			}
		}
	}
}

func extractAbnormalFlow(r *core.AggreSpanRecord) ([]string, map[string][]string, map[string][]string) {
	var keys []string
	lossFlowMap := make(map[string][]string)
	bitFlipFlowMap := make(map[string][]string)
	for k, v := range r.AggreRecordMap {
		var lossFlows []string
		var bitFlipFlows []string
		lossMap := make(map[string]bool)
		bitFlipMap := make(map[string]bool)
		for _, pktRecord := range v.PktRecords {
			var key string
			if pktRecord.Proto == "icmp" {
				key = fmt.Sprintf("%s->%s", pktRecord.SrcIP, pktRecord.DstIP)
			} else {
				key = fmt.Sprintf("%s:%d->%s:%d", pktRecord.SrcIP, pktRecord.SrcPort, pktRecord.DstIP, pktRecord.DstPort)
			}
			if !pktRecord.Received && !lossMap[key] {
				lossMap[key] = true
				lossFlows = append(lossFlows, key)
			}

			if pktRecord.BitFlip && !bitFlipMap[key] {
				bitFlipMap[key] = true
				bitFlipFlows = append(bitFlipFlows, key)
			}
		}

		if len(lossFlows) > 0 {
			sort.Strings(lossFlows)
			lossFlowMap[k] = lossFlows
		}

		if len(bitFlipFlows) > 0 {
			sort.Strings(bitFlipFlows)
			bitFlipFlowMap[k] = bitFlipFlows
		}

		if len(lossFlowMap) > 0 || len(bitFlipFlowMap) > 0 {
			keys = append(keys, k)
		}
	}

	return keys, lossFlowMap, bitFlipFlowMap
}
